#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <liburing.h>
#include <coroutine>
#include <thread> // for std::terminate

#ifndef BLOCK_SZ
// You can define BLOCK_SZ if you want to. This is not documented.
#define BLOCK_SZ 1024
#endif

struct file_info
{
    off_t file_sz;
    iovec iovecs[]; // Referred by readv/writev
};

//  Returns the size of the file whose open file descriptor is passed in.
//  Properly handles regular file and block devices as well. Pretty.
off_t get_file_size(int fd)
{
    struct stat st;

    if (fstat(fd, &st) < 0)
    {
        perror("fstat");
        return -1;
    }
    if (S_ISBLK(st.st_mode))
    {
        unsigned long long bytes;
        if (ioctl(fd, BLKGETSIZE64, &bytes) != 0)
        {
            perror("ioctl");
            return -1;
        }
        return bytes;
    }
    else if (S_ISREG(st.st_mode))
        return st.st_size;

    return -1;
}

// Submit the readv request via liburing
int create_read_request(int file_fd, off_t offset, size_t length, char *buf, io_uring *ring, off_t &f_sz)
{
    if (file_fd < 0)
    {
        perror("open");
        return 1;
    }

    if (offset < 0) [[unlikely]]
        offset = 0;

    off_t file_sz = get_file_size(file_fd);
    if (offset > file_sz)
    {
        fprintf(stderr, "Offset larger than file size.\n");
        return 1;
    }
    size_t bytes_remaining = std::min(size_t(file_sz - offset), length);
    f_sz = bytes_remaining;

    int current_block = 0;
    int blocks = bytes_remaining / BLOCK_SZ;
    if ((bytes_remaining % BLOCK_SZ) && bytes_remaining)
        ++blocks;

    file_info *fi = (file_info *)malloc(sizeof(file_info) + sizeof(iovec) * blocks);
    fi->file_sz = bytes_remaining;

    //   For each block of the file we need to read, we allocate an iovec struct
    //   which is indexed into the iovecs array. This array is passed in as part
    //   of the submission. If you don't understand this, then you need to look
    //   up how the readv() and writev() system calls work.
    unsigned long long index = 0;

    while (bytes_remaining)
    {
        off_t bytes_to_read = bytes_remaining;
        if (bytes_to_read > BLOCK_SZ)
            bytes_to_read = BLOCK_SZ;

        fi->iovecs[current_block].iov_len = bytes_to_read;
        fi->iovecs[current_block].iov_base = (char *)((unsigned long long)buf + index);

        current_block++;
        bytes_remaining -= bytes_to_read;
        index += bytes_to_read;
    }

    // Get an SQE
    io_uring_sqe *sqe = io_uring_get_sqe(ring);
    // Setup a readv operation
    io_uring_prep_readv(sqe, file_fd, fi->iovecs, blocks, offset);
    // Set user data
    io_uring_sqe_set_data(sqe, fi);

    return 0;
}

class io_uring_wrapper
{
public:
    struct promise_type;
    using handle = std::coroutine_handle<promise_type>;
    struct promise_type
    {
        promise_type(char *filename, off_t offset, size_t buffer_length, char *buf, io_uring *ring) : _ring(ring)
        {
            if (!ring)
            {
                fprintf(stderr, "Must provide ring.");
                throw;
            }
            // Open file
            _file_fd = open(filename, O_RDONLY);
            int ret = create_read_request(_file_fd, offset, buffer_length, buf, ring, this->_file_sz);
            if (ret)
            {
                fprintf(stderr, "Error reading file: %s\n", filename);
                throw;
            }
        }

        promise_type(int file_fd, off_t offset, size_t buffer_length, char *buf, io_uring *ring) : _ring(ring)
        {
            if (!ring)
            {
                fprintf(stderr, "Must provide ring.");
                throw;
            }
            int ret = create_read_request(file_fd, offset, buffer_length, buf, ring, this->_file_sz);
            if (ret)
            {
                fprintf(stderr, "Error using file descriptor: %d\n", file_fd);
                throw;
            }
        }

        io_uring_wrapper get_return_object() { return io_uring_wrapper{handle::from_promise(*this)}; }
        static io_uring_wrapper get_return_object_on_allocation_failure() noexcept { return io_uring_wrapper{nullptr}; }

        struct final_awaiter
        { // std::suspend_never is not nothrow-constructible, while final awaiter must be. A work-around here.
            bool await_ready() noexcept { return true; }
            void await_suspend(handle h) noexcept {}
            void await_resume() noexcept {}
        };

        auto initial_suspend() { return std::suspend_always{}; }
        auto final_suspend() noexcept { return final_awaiter{}; }
        void unhandled_exception() { std::terminate(); }
        void return_void() noexcept
        {
            // Call the clean-up function.
            if (_file_fd != -1)
                close(_file_fd);
        }
        auto yield_value(int)
        {
            io_uring_cqe *cqe;
            int ret = io_uring_wait_cqe(_ring, &cqe);
            if (ret < 0)
            {
                perror("io_uring_wait_cqe");
                throw;
            }
            if (cqe->res < 0)
            {
                fprintf(stderr, "Async readv failed.\n");
                throw;
            }
            free(io_uring_cqe_get_data(cqe));
            io_uring_cqe_seen(_ring, cqe);
            return std::suspend_always{};
        }
        io_uring *_ring;
        off_t _file_sz = -1;
        int _file_fd = -1;
    };

    io_uring_wrapper(const io_uring_wrapper &) = delete;
    io_uring_wrapper(io_uring_wrapper &&) = default;
    io_uring_wrapper &operator=(const io_uring_wrapper &) = delete;
    io_uring_wrapper &operator=(io_uring_wrapper &&) = default;

    size_t operator()()
    {
        switch (++state)
        {
        case 1:
            // wait until ready
            coro.resume();
            return coro.promise()._file_sz;

        case 2:
            // release file
            coro.resume();
            // coro is also destroyed here. Calling coro after here would cause a seg fault.
            return 0;

        default:
            return -1;
        }
    }

private:
    io_uring_wrapper(handle h) : coro(h) {}
    handle coro;
    int state = 0;

public:
    ~io_uring_wrapper() noexcept
    {
        // ensure the coro is resumed exactly twice
        if (coro && state == 1) [[unlikely]]
        {
            coro.resume();
        }
    }
};

// by filename
io_uring_wrapper request_file_by_name(char *filename, off_t offset, size_t buffer_length, char *buf, io_uring *ring)
{
    co_yield 0;
    co_return;
}

// by file descriptor
io_uring_wrapper request_file_by_fd(int file_fd, off_t offset, size_t buffer_length, char *buf, io_uring *ring)
{
    co_yield 0;
    co_return;
}
